import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import os

class UltrasoundClassifier:
    def __init__(self, device=None, labels=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        # 使用 EfficientNet-B0 作为骨干，采用 torchvision 的预训练权重（ImageNet）
        try:
            # For torchvision >= 0.13, weights argument exists; if not, this falls back to default
            self.model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1)
        except Exception:
            self.model = models.efficientnet_b0(pretrained=True)
        # 修改分类头以匹配类别数
        in_features = self.model.classifier[1].in_features
        num_classes = 2
        self.model.classifier[1] = nn.Linear(in_features, num_classes)
        self.model.to(self.device)
        self.model.eval()

        # 标签映射（可按需扩展）
        self.labels = labels or ["liver", "thyroid"]

        # 预处理
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ])

        # 如果存在微调权重（可选），自动加载
        possible_weight = os.path.join(os.path.dirname(__file__), "..", "weights", "classifier_imagenet.pth")
        if os.path.exists(possible_weight):
            try:
                state = torch.load(possible_weight, map_location=self.device)
                self.model.load_state_dict(state)
                print("[Classifier] Loaded fine-tuned weights from:", possible_weight)
            except Exception as e:
                print("[Classifier] Failed to load fine-tuned weights:", e)

    def predict(self, image_path: str) -> str:
        image = Image.open(image_path).convert("RGB")
        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
        with torch.no_grad():
            outputs = self.model(img_tensor)
            probs = torch.softmax(outputs, dim=1)
            _, predicted = torch.max(probs, 1)
        return self.labels[predicted.item()]
